#!/usr/bin/env python
# -*- coding: utf-8 -*-

import datetime
import decimal
from contextlib import contextmanager

from flask_sqlalchemy import SQLAlchemy, BaseQuery as _BaseQuery

from fx.lib.utils import FancyDict
from fx.utils.misc import id_encrypt, logger


__all__ = [
    'db',
    'BaseQuery',
    'SessionMixin',
    'DefaultModel',
    'BaseModel',
    'session_scope'
]

db = SQLAlchemy()


class BaseQuery(_BaseQuery):

    def filter_in(self, model, values):
        values = set(values)
        if len(values) == 0:
            return {}
        if len(values) == 1:
            ident = values.pop()
            rv = self.get(ident)
            if not rv:
                return {}
            return {ident: rv}
        items = self.filter(model.in_(values))
        dct = {}
        for item in items:
            dct[getattr(item, model.key)] = item
        return dct

    def as_list(self):
        return list(self)


db.Model.query_class = BaseQuery


class SessionMixin(object):

    def save(self, **kwargs):
        try:
            db.session.add(self)
            db.session.commit()
            return self
        except Exception as e:
            logger.error(e, exc_info=True)
            db.session.rollback()
            raise

    def delete(self, **kwargs):
        try:
            db.session.delete(self)
            db.session.commit()
            return self
        except Exception as e:
            logger.error(e, exc_info=True)
            db.session.rollback()
            raise


class DefaultModel(db.Model):
    __abstract__ = True
    __table_args__ = {'mysql_charset': 'utf8mb4', 'mysql_engine': 'InnoDB'}


class BaseModel(DefaultModel, SessionMixin):
    __abstract__ = True

    dict_default_columns = []

    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            setattr(self, k, v)

    def to_dict(self, columns=None):
        dct = FancyDict()
        if columns is None:
            columns = self.dict_default_columns
        for col in columns:
            if (col == 'id' or col in getattr(self, 'encrypt_attrs', [])) and getattr(self, col):
                value = id_encrypt(getattr(self, col))
            else:
                value = getattr(self, col)
            if isinstance(value, datetime.datetime):
                value = value.strftime('%Y-%m-%d %H:%M:%S')
            if isinstance(value, decimal.Decimal):
                value = str(value)
            if isinstance(value, BaseModel):
                value = value.to_dict()
            dct[col] = value
        return dct


@contextmanager
def session_scope():
    session = db.session
    try:
        yield session
        session.commit()
    except:
        session.rollback()
        raise
